#include "vc.h"

int Vcycle_1d
(
 const double *v_in, 
 const double *f_in, 
 int size_in,  
 int nu_down, 
 int nu_up,
 double *v_out
)
{
  int i,j,k;
  int N = size_in - 1 ;
  int num_level = log( N ) / log(2);
  int len_arr = N+1;
  int len_arr_vec[num_level];
  int size_vec[num_level];
  int cur_N = N;
  double h = 0.5;
  double hsq = h * h;

  len_arr_vec[0] = N+1;
  size_vec[0] = N+1;

  for(i = 1; i < num_level; i++)
    {
      cur_N = cur_N / 2;
      len_arr += (cur_N+1);
      len_arr_vec[i] = len_arr;
      size_vec[i] = cur_N + 1;
      //      printf("len_arr is %d \n", len_arr);
      //      printf("num_grd is %d \n", size_vec[i]);
    }
  
  double *v = (double *) malloc (len_arr * sizeof(double));
  double *f = (double *) malloc (len_arr * sizeof(double));
  // pointer to the beginning of the v array
  double **v_ptr_vec = (double **) malloc(num_level * sizeof(double *));
  // pointer to the beginning of the f array
  double **f_ptr_vec = (double **) malloc(num_level * sizeof(double *)); 
  double *v_ptr = v;
  double *f_ptr = f;

  for(i=0; i<len_arr; i++)
    {
      v[i] = 0.0;
      f[i] = 0.0;
    }

  for(i = 0; i < num_level; i++)
    {
      v_ptr_vec[i] = v_ptr ;
      f_ptr_vec[i] = f_ptr ;
      v_ptr += 	size_vec[i];
      f_ptr +=  size_vec[i];
    }
  //  double *ie = v + sizeof(double) * ; // pointer to the end of the sub array

  // Copy the v_in and f_in to the beginning of the long vectors
  for(i = 0; i < len_arr_vec[0]; i++)
    {
      v[i] = v_in[i];
      f[i] = f_in[i];
    }
  
  // Begin of V-Cycle
  // Traverse down grids
  for(j = 0; j < num_level-2; j++)  // j < num_level -1 if not using serial jacabi
    {
      relax(f_ptr_vec[j], size_vec[j], nu_down, v_ptr_vec[j]);
      restrict(v_ptr_vec[j], f_ptr_vec[j], size_vec[j], f_ptr_vec[j+1] );
    }

  // solve exactly on coarsest grid (3 points)
  /*
  *(v_ptr_vec[num_level-1]) = 0;
  *(v_ptr_vec[num_level-1] + 1) = (hsq / 2) * ( *(f_ptr_vec[num_level-1] + 1) ) ;
  *(v_ptr_vec[num_level-1] + 2) = 0;
  */
  serial_jacobi(f_ptr_vec[num_level-2], size_vec[num_level-2], v_ptr_vec[num_level-2]);
					       
  /* Traverse up grids */
  for(j = num_level - 3; j >= 0; j--)  //j=num_level -2 if not using serial jocobi
    {
      interpolate(v_ptr_vec[j+1], size_vec[j+1], v_ptr_vec[j]);  //becareful with the indices
      relax(f_ptr_vec[j], size_vec[j], nu_up, v_ptr_vec[j]);
    } 

  /* Copy the solution to the output vector */
  for(i = 0; i < size_in; i++)
    v_out[i] = v[i];

  /* Print out the solution */
  /*
  printf("index:    f     v  \n-------------\n");
  for(j = 0; j < size_vec[0]; j++)
    printf("%3i    %10.4f   %10.4f \n", j, f[j], v[j]);
  */
  
  free(v);
  free(f);
  free(v_ptr_vec);
  free(f_ptr_vec);

  return 0;
}


int interpolate
(
 double *v_in,
 int size_in,
 double *v_out
)
{

  int nc = size_in - 1;   // number of coase grids
  int nf = 2 * nc;        // number of fine grids
  double v_delta[nf+1];

  int i;
  // OMP
  for(i=2; i < nf-1; i += 2 )  //for odd indices copy the values
    v_delta[i] = v_in[i/2];

  // OMP
  for(i=1; i < nf; i += 2 )    //for even indices interpolate
    v_delta[i] = 0.5 * ( v_in[ (i-1)/2] + v_in[ (i-1)/2 + 1] );
       
  v_delta[0] = 0;
  v_delta[nf] = 0;

  for(i=0; i < nf+1; i++)
    v_out[i] += v_delta[i];

  for(i=0; i < nc+1; i++)
    v_in[i] = 0;

  return 0;
}

int restrict
(
 const double *v_in,
 const double *f_in,
 int size_in,
 double *f_out
 )
{
  int i,j;
  int nf = size_in - 1;   // number of fine grids
  int nc = nf / 2;        // number of coarse grids

  double *f_delta = (double*) malloc(sizeof(double) * (nf+1));
  double *v_tmp   = (double*) malloc(sizeof(double) * (nf+1));

  apply(v_in, size_in, v_tmp);

  // Get the residual
  for(i = 0; i < nf+1; i++)
    f_delta[i] = f_in[i] - v_tmp[i];

  f_out[0] = 0;
  f_out[nc] = 0;

  //printf("now nc is %d, nf is %d\n", nc, nf);

  for(i = 1; i < nc; i++)
    f_out[i] = ( f_delta[2*i-1] + 2*f_delta[2*i] + f_delta[2*i+1] ) / 4;

  free(f_delta);
  free(v_tmp);
  return 0;
}

int apply
(
 const double *v_in,
 int size,
 double *v_out
 )
{
  int n = size-1;
  double h = 1.0 / n;
  double hsq = h*h;
  int i;

  v_out[0] = 0;
  v_out[n] = 0;

  // OMP  
  for(i = 1; i < n; i++)
    v_out[i] = 1/hsq * (-v_in[i-1] + 2*v_in[i] - v_in[i+1]);

  return 0;
}


int relax
(
 const double *f_in,
 int size_in,
 int nu,
 double *v_in_out
 )
{
  double w = 2.0/3.0;
  int N = size_in - 1;
  double h = 1.0 / N; 
  double hsq = h * h;
  int i, j;
  double *v_tmp = (double *) malloc(sizeof(double) * size_in);
  // OMP + MPI

  for(i = 0; i < nu; i++)
    {

#pragma omp parallel shared(v_tmp, v_in_out, f_in, N, w, hsq) 
      {
#pragma omp for private(j) 
	for(j = 1; j < N; j++)
	  {
	    //	  v_in_out[j] = (1-w) * v_in_out[j] + (w/2) * (v_in_out[j-1] + v_in_out[j+1] + hsq * f_in[j]);
	    v_tmp[j] = (1-w) * v_in_out[j] + (w/2) * (v_in_out[j-1] + v_in_out[j+1] + hsq * f_in[j]);
	  }

#pragma omp single
	{      
	  v_tmp[0] = 0;
	  v_tmp[N] = 0;
	}

#pragma omp for private(j)
	for(j = 0; j < N+1; j++)
	  v_in_out[j] = v_tmp[j];      

      }

      //      v_in_out[0] = 0;
      //      v_in_out[N] = 0;

      if(VERBOSE == 1)
	{
	  printf("\n\niteration %i \n -------------\n", i);
	  for(j = 0; j < N+1; j++)
	    printf("v[%i]: %f \n", j, v_in_out[j]);
	}
    }

  free(v_tmp);
  return 0;
}


int serial_jacobi
(
 const double *f_in,
 int size,
 double *v_in_out
)
{

  double v_tmp[size];
  int i,j;
  double error;
  double tol = 1e-6;
  double h = 1.0 / (size-1); 
  double hsq = h * h;
  int flag = 0;
  int iter = 0;

  /*
  for(i = 0; i < size; i++)
    {
      printf("v_in_out[%i]: %f of size %d\n", i, v_in_out[i], size);
    }
  */

  while(flag <= 0 && iter < 500)
    {
      iter += 1;

      for(j = 1; j < size-1; j++)
	v_tmp[j] = 0.5 * (v_in_out[j-1] + v_in_out[j+1] + hsq * f_in[j]);
      v_tmp[0] = 0;
      v_tmp[size-1] = 0;
      /*
      printf("\n");
      for(i=0; i < size; i++) 
	{
	  v_in_out[i] = v_tmp[i];
	  printf("v_in_out[%i]: %f\n", i, v_in_out[i]);
	}
      printf("\n");
      */
      for(i=0; i < size; i++)
	{
	  //	  printf("v_in_out[%i]: %f\n", i, v_in_out[i]);
	  error = fabs(v_in_out[i] - f_in[i]);
	  //	  printf("error is %f \n", error);
	  if(error > tol)
	    {
	      break;
	    }
	  if(i == size -1) flag = 1;
	}

    }

  //  printf("Serial Jacobi takes %i iterations with maximum error %f\n", 
  //	 iter, error);

  return 0;

}


double get_2norm_error(double* v, int size)
{
  double esq = 0, e = 0;
  int i;
  for(i = 0; i < size; i++)
    esq += v[i]*v[i];
  e = sqrt(esq);
  return e;
}
